"""
查找最合适的分割位置
"""
from parse_config import Config
import json

def load_json(json_path):
    with open(json_path,'r') as f:
        resource_nodes=json.load(f)
    return resource_nodes

def find_greatest_split(resource_nodes,config):
    gpu_memory=config['gpu_memory'] # 显存
    gpu_speed=config['gpu_speed'] # gpu和cpu的速度比值
    network_speed=config['network_speed'] # 在不同设备之间传输的带宽
    cpu_nodes=[]
    result_list = [{"layer": k, **v} for k, v in resource_nodes.items()]
    for i in range(len(result_list)):
        if result_list[i]['peak_memory']>=gpu_memory*1024:
            cpu_nodes.append(i)
    print([result_list[i]['peak_memory'] for i in range(len(result_list))])
    if len(cpu_nodes)==0:
        print('可以让完整模型在gpu上运行')
        return []
    else:
        print('模型需要部分转移到cpu上运行')
        ans_nodes=set()
        for i in range(len(cpu_nodes)):
            cur_node=cpu_nodes[i]
            if i==0:
                time_cost=0
                j=0
                while j < cur_node:
                    time_cost+=result_list[j]['time_cost']*(gpu_speed-1)
                    j+=1
                if time_cost<=result_list[cur_node-1]['output_memory']/(network_speed*1024):
                    for s in range(cur_node+1):
                        ans_nodes.add(s)
                else:
                    ans_nodes.add(cur_node)
            else:
                time_cost=0
                j=cpu_nodes[i-1]+1
                while j < cur_node:
                    time_cost+=result_list[j]['time_cost']*(gpu_speed-1)
                    j+=1
                if time_cost<=(result_list[cur_node-1]['output_memory']+result_list[cpu_nodes[i-1]]['output_memory'])/(network_speed*1024):
                    for s in range(cpu_nodes[i-1]+1,cur_node+1):
                        ans_nodes.add(s)
                else:
                    ans_nodes.add(cur_node)
        return list(ans_nodes), [result['layer'] for result in result_list]

def main():
    config_path='./config.json'
    config=Config(config_path)
    resource_nodes=load_json(config.config['valid_resourse_cost_save_path'])
    cpu_nodes=find_greatest_split(resource_nodes,config.config)
    print(cpu_nodes)
    return cpu_nodes
    
    
if __name__ == '__main__':
    main()